package com.wxl.asmgradleplugin.myplugin

import javassist.ClassPool
import org.gradle.api.artifacts.transform.InputArtifact
import org.gradle.api.artifacts.transform.TransformAction
import org.gradle.api.artifacts.transform.TransformOutputs
import org.gradle.api.artifacts.transform.TransformParameters
import org.gradle.api.file.FileSystemLocation
import org.gradle.api.provider.Provider
import java.io.FileOutputStream
import java.util.jar.JarEntry
import java.util.jar.JarFile
import java.util.jar.JarOutputStream

abstract class JarInstrumentTransform : TransformAction<TransformParameters.None> {

    @get:InputArtifact
    abstract val inputArtifact: Provider<FileSystemLocation>

    override fun transform(outputs: TransformOutputs) {
        val inputJar = inputArtifact.get().asFile
        println("===444 input ${inputJar.absolutePath}")
        val outputJar = outputs.file(inputJar.name)

        val classPool = ClassPool.getDefault()
        val jarFile = JarFile(inputJar)
        val jarOut = JarOutputStream(FileOutputStream(outputJar))

        for (entry in jarFile.entries()) {
            val inputStream = jarFile.getInputStream(entry)
            val newEntry = JarEntry(entry.name)
            jarOut.putNextEntry(newEntry)

            if (entry.name.endsWith(".class")) {
                val originalBytes = inputStream.readBytes()

                try {
                    val ctClass = classPool.makeClass(originalBytes.inputStream())

                    if (!ctClass.isFrozen && !ctClass.isInterface) {
                        ctClass.defrost() // 有时候 Gradle 会缓存冻结类
                        ctClass.constructors.forEach {
                            it.insertBefore("""System.out.println(">>> JAR Constructor called");""")
                        }
                        val modifiedBytes = ctClass.toBytecode()
                        jarOut.write(modifiedBytes)
                        ctClass.detach()
                    } else {
                        jarOut.write(originalBytes)
                    }
                } catch (e: Exception) {
                    // 某些类可能无法修改，比如 module-info.class，写回原始内容
                    jarOut.write(originalBytes)
                }
            } else {
                // 非 class 文件直接复制
                jarOut.write(inputStream.readBytes())
            }

            jarOut.closeEntry()
        }

        jarOut.close()
        jarFile.close()
    }
}
